# _*_ coding : UTF-8 _*_
# 开发人员   : TestLin
# 开发时间   : 2024/6/6  10:13
# 文件名称   : app.py
# 开发工具   : PyCharm

import itchat
from itchat.content import TEXT, PICTURE, MAP, CARD, NOTE, SHARING, RECORDING, ATTACHMENT, VIDEO
from flask import Flask, request, jsonify

app = Flask(__name__)

# Predefined key for validation
API_KEY = "1234567890"


class WeChatAPI:
    def __init__(self):
        self.msg_callbacks = []

    def login(self):
        itchat.auto_login(enableCmdQR=2, hotReload=True)
        #itchat.auto_login(hotReload=True)
        return "Logged in successfully."

    def logout(self):
        itchat.logout()
        return "Logged out successfully."

    def send_message(self, to_user, message, is_group=False):
        if is_group:
            chatroom = itchat.search_chatrooms(name=to_user)
            if chatroom:
                to_user = chatroom[0]['UserName']
            else:
                return f"Group '{to_user}' not found."
        else:
            friend = itchat.search_friends(nickName=to_user) or \
                     itchat.search_friends(remarkName=to_user) or \
                     itchat.search_friends(alias=to_user)
            if friend:
                to_user = friend[0]['UserName']
            else:
                return f"Friend '{to_user}' not found."
        itchat.send(message, toUserName=to_user)
        return "Message sent successfully."

    def get_friend_list(self):
        friends = itchat.get_friends(update=True)[1:]  # Skip the first one (self)
        friend_list = [
            {'nickName': friend['NickName'], 'remarkName': friend.get('RemarkName'), 'alias': friend.get('Alias'),
             'userName': friend['UserName']} for friend in friends]
        return friend_list

    def get_group_list(self):
        chatrooms = itchat.get_chatrooms(update=True)
        group_list = [{'name': room['NickName'], 'userName': room['UserName']} for room in chatrooms]
        return group_list

    def add_msg_callback(self, callback):
        self.msg_callbacks.append(callback)

    def msg_dispatcher(self, msg):
        for callback in self.msg_callbacks:
            callback(msg)

    def run(self):
        @itchat.msg_register([TEXT, PICTURE, MAP, CARD, NOTE, SHARING, RECORDING, ATTACHMENT, VIDEO])
        def handle_receive_msg(msg):
            self.msg_dispatcher(msg)

        itchat.run()


wechat_api = WeChatAPI()


def verify_key(data):
    provided_key = data.get('key')
    return provided_key == API_KEY


@app.route('/login', methods=['POST'])
def login():
    if not request.is_json or not verify_key(request.json):
        return jsonify({"error": "Unauthorized"}), 401
    return wechat_api.login()


@app.route('/logout', methods=['POST'])
def logout():
    if not request.is_json or not verify_key(request.json):
        return jsonify({"error": "Unauthorized"}), 401
    return wechat_api.logout()


@app.route('/send_message', methods=['POST'])
def send_message():
    if not request.is_json or not verify_key(request.json):
        return jsonify({"error": "Unauthorized"}), 401

    data = request.json
    to_user = data.get('to_user')
    message = data.get('message')
    is_group = data.get('is_group', False)

    if not to_user or not message:
        return jsonify({"error": "Missing 'to_user' or 'message'"}), 400

    return wechat_api.send_message(to_user, message, is_group)


@app.route('/get_friends', methods=['GET'])
def get_friends():
    if request.args.get('key') != API_KEY:
        return jsonify({"error": "Unauthorized"}), 401

    try:
        friend_list = wechat_api.get_friend_list()
        return jsonify(friend_list), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/get_groups', methods=['GET'])
def get_groups():
    if not request.args.get('key') == API_KEY:
        return jsonify({"error": "Unauthorized"}), 401
    group_list = wechat_api.get_group_list()
    return jsonify(group_list)


def print_msg(msg):
    print(f"Received message: {msg['Text']} from {msg['FromUserName']}")


wechat_api.add_msg_callback(print_msg)

if __name__ == "__main__":
    wechat_api.login()
    print("Logged in successfully")
    # Run the wechat message handler in a separate thread
    import threading

    threading.Thread(target=wechat_api.run).start()
    # Run the Flask app
    app.run(host='0.0.0.0', port=5000)